import numpy as np
import matplotlib.pyplot as plt
import scipy.io
import random
import matplotlib
import MDSplus
import os

# 设置matplotlib支持中文显示
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 用黑体显示中文
matplotlib.rcParams['axes.unicode_minus'] = False    # 正常显示负号


def kalman_unwrap(phase, phase_range=8, process_noise=0.01, measurement_noise=0.1):
    """
    卡尔曼滤波器解缠绕
    使用卡尔曼滤波器预测相位变化，更适合噪声较大的情况
    
    参数:
    - phase: 相位信号
    - phase_range: 相位范围，默认为8 (-4到4)
    """
    # 将相位范围从[-4,4]归一化到[-π,π]
    phase_normalized = phase * np.pi / (phase_range/2)
    
    unwrapped = np.zeros_like(phase_normalized)
    unwrapped[0] = phase_normalized[0]
    
    # 卡尔曼滤波器状态
    x_est = phase_normalized[0]  # 状态估计
    p_est = 1.0       # 估计误差协方差
    
    for i in range(1, len(phase_normalized)):
        # 预测步骤
        x_pred = x_est
        p_pred = p_est + process_noise
        
        # 计算相位差
        diff = phase_normalized[i] - (x_pred % (2 * np.pi))
        
        # 调整差值到[-π,π]范围
        if diff > np.pi:
            diff -= 2 * np.pi
        elif diff < -np.pi:
            diff += 2 * np.pi
        
        # 更新步骤
        k = p_pred / (p_pred + measurement_noise)  # 卡尔曼增益
        x_est = x_pred + k * diff
        p_est = (1 - k) * p_pred
        
        unwrapped[i] = x_est
    
    # 将结果转换回原始相位范围
    return unwrapped * (phase_range/2) / np.pi


def process_density(phase_signal, phase_range=8, density_factor=1.0):
    """
    将相位信号处理为密度信号
    
    参数:
    - phase_signal: 相位信号
    - phase_range: 相位范围
    - density_factor: 密度计算的校正系数
    
    返回:
    - density: 处理后的ne
    """
    # 使用卡尔曼滤波器解缠绕
    unwrapped_phase = kalman_unwrap(phase_signal, phase_range)
    
    # 使用指定公式计算密度
    dl = 0
    density = (unwrapped_phase - dl) / 4 * np.pi * 35.5 / 3.37 / 2 * density_factor
    
    return density


def plot_comparison(file_name='sjj1.mat', num_samples=4, density_factor=1.0):
    """
    绘制卡尔曼方法处理后的数据与真实数据的对比图
    
    参数:
    - file_name: 数据文件名
    - num_samples: 要绘制的样本数量
    - density_factor: 密度计算的校正系数
    """
    # 加载数据
    data = scipy.io.loadmat(file_name)
    
    # 随机抽取样本
    indices = random.sample(range(1, 51), num_samples)
    
    # 创建子图
    fig, axes = plt.subplots(num_samples, 1, figsize=(12, 4*num_samples))
    
    # 如果只有一个样本，确保axes是列表
    if num_samples == 1:
        axes = [axes]
    
    # 处理每个样本
    for i, idx in enumerate(indices):
        # 获取原始相位和真实密度
        y_original = data[f'fai{idx}'].flatten()
        y_true = data[f'ne{idx}'].flatten()
        
        # 确保长度一致
        min_length = min(len(y_original), len(y_true))
        y_original = y_original[:min_length]
        y_true = y_true[:min_length]
        
        # 处理相位信号
        y_processed = process_density(y_original, density_factor=density_factor)
        
        # 计算均方根误差和相关系数
        mse = np.mean((y_processed - y_true) ** 2)
        corr = np.corrcoef(y_processed, y_true)[0, 1]
        
        # 绘制对比图
        axes[i].plot(y_processed, label='卡尔曼处理后')
        axes[i].plot(y_true, 'r--', label='真实密度')
        axes[i].set_title(f'样本 {idx}: MSE = {mse:.4e}, 相关系数 = {corr:.4f}')
        axes[i].set_xlabel('时间点')
        axes[i].set_ylabel('ne')
        axes[i].legend()
        axes[i].grid(True)
    
    plt.tight_layout()
    plt.savefig(f'e:\\ENN\\2025工作\\诊断-hcn拼接\\kalman_comparison.png')
    plt.show()


def optimize_parameters(file_name='sjj1.mat', test_indices=None):
    """
    优化卡尔曼滤波器参数
    
    参数:
    - file_name: 数据文件名
    - test_indices: 用于测试的样本索引，如果为None则随机选择
    
    返回:
    - best_params: 最佳参数组合
    """
    # 加载数据
    data = scipy.io.loadmat(file_name)
    
    # 如果未指定测试索引，随机选择
    if test_indices is None:
        test_indices = random.sample(range(1, 51), 3)
    
    # 参数网格
    process_noise_values = [0.001, 0.01, 0.1, 1.0]
    measurement_noise_values = [0.01, 0.1, 1.0, 10.0]
    density_factor_values = [0.8, 0.9, 1.0, 1.1, 1.2]
    
    best_mse = float('inf')
    best_params = {}
    
    # 网格搜索
    for pn in process_noise_values:
        for mn in measurement_noise_values:
            for df in density_factor_values:
                total_mse = 0
                
                for idx in test_indices:
                    # 获取原始相位和真实密度
                    y_original = data[f'fai{idx}'].flatten()
                    y_true = data[f'ne{idx}'].flatten()
                    
                    # 确保长度一致
                    min_length = min(len(y_original), len(y_true))
                    y_original = y_original[:min_length]
                    y_true = y_true[:min_length]
                    
                    # 自定义卡尔曼滤波器参数
                    unwrapped = kalman_unwrap(y_original, process_noise=pn, measurement_noise=mn)
                    
                    # 使用指定公式计算密度
                    dl = 0
                    density = (unwrapped - dl) * 35.5 / 3.37 / 2 * df
                    
                    # 计算MSE
                    mse = np.mean((density - y_true) ** 2)
                    total_mse += mse
                
                # 计算平均MSE
                avg_mse = total_mse / len(test_indices)
                
                # 更新最佳参数
                if avg_mse < best_mse:
                    best_mse = avg_mse
                    best_params = {
                        'process_noise': pn,
                        'measurement_noise': mn,
                        'density_factor': df
                    }
    
    print(f"最佳参数: {best_params}, 平均MSE: {best_mse:.4e}")
    return best_params


def plot_single_sample_detail(file_name, sample_idx, process_noise, measurement_noise, density_factor):
    """
    绘制单个样本的详细对比图
    
    参数:
    - file_name: 数据文件名
    - sample_idx: 样本索引
    - process_noise: 卡尔曼滤波器过程噪声参数
    - measurement_noise: 卡尔曼滤波器测量噪声参数
    - density_factor: 密度计算的校正系数
    
    返回:
    - mse: 均方根误差
    - corr: 相关系数
    """
    # 加载数据
    data = scipy.io.loadmat(file_name)
    
    # 获取原始相位和真实密度
    y_original = data[f'fai{sample_idx}'].flatten()
    y_true = data[f'ne{sample_idx}'].flatten()
    
    # 确保长度一致
    min_length = min(len(y_original), len(y_true))
    y_original = y_original[:min_length]
    y_true = y_true[:min_length]
    
    # 使用指定参数处理
    unwrapped = kalman_unwrap(
        y_original, 
        process_noise=process_noise, 
        measurement_noise=measurement_noise
    )
    
    # 使用指定公式计算密度
    dl = 0
    y_processed = (unwrapped - dl) * 35.5 / 3.37 / 2 * density_factor
    
    # 计算评估指标
    mse = np.mean((y_processed - y_true) ** 2)
    corr = np.corrcoef(y_processed, y_true)[0, 1]
    
    # 绘制详细对比图
    plt.figure(figsize=(12, 8))
    
    # 上图：原始相位信号和解缠绕后的相位
    plt.subplot(2, 1, 1)
    plt.plot(y_original, label='原始相位信号')
    plt.plot(unwrapped, label='解缠绕后的相位')
    plt.title('相位信号对比')
    plt.xlabel('时间点')
    plt.ylabel('相位')
    plt.legend()
    plt.grid(True)
    
    # 下图：处理后的密度和真实密度
    plt.subplot(2, 1, 2)
    plt.plot(y_processed, label='卡尔曼处理后')
    plt.plot(y_true, 'r--', label='真实密度')
    plt.title(f'样本 {sample_idx} 密度对比: MSE = {mse:.4e}, 相关系数 = {corr:.4f}')
    plt.xlabel('时间点')
    plt.ylabel('ne')
    plt.legend()
    plt.grid(True)
    
    plt.tight_layout()
    
    # 创建结果目录
    result_dir = 'e:\\ENN\\2025工作\\诊断-hcn拼接\\result'
    import os
    if not os.path.exists(result_dir):
        os.makedirs(result_dir)
    
    # 保存图像
    plt.savefig(f'{result_dir}\\kalman_sample_{sample_idx}_detail.png')
    plt.close()  # 关闭图形，避免内存泄漏
    
    return mse, corr


def plot_all_samples(file_name, best_params):
    """
    绘制所有样本的详细对比图
    
    参数:
    - file_name: 数据文件名
    - best_params: 最佳参数字典
    """
    # 创建结果目录
    result_dir = 'e:\\ENN\\2025工作\\诊断-hcn拼接\\result'
    import os
    if not os.path.exists(result_dir):
        os.makedirs(result_dir)
    
    # 准备存储所有样本的评估结果
    results = []
    
    # 循环处理所有样本
    for idx in range(1, 51):
        try:
            print(f"处理样本 {idx}...")
            mse, corr = plot_single_sample_detail(
                file_name,
                idx,
                best_params['process_noise'],
                best_params['measurement_noise'],
                best_params['density_factor']
            )
            results.append({
                'sample_idx': idx,
                'mse': mse,
                'correlation': corr
            })
        except Exception as e:
            print(f"处理样本 {idx} 时出错: {e}")
    
    # 生成汇总报告
    with open(f'{result_dir}\\summary_report.txt', 'w', encoding='utf-8') as f:
        f.write("卡尔曼滤波器相位解缠绕结果汇总\n")
        f.write("================================\n\n")
        f.write(f"使用参数:\n")
        f.write(f"  过程噪声: {best_params['process_noise']}\n")
        f.write(f"  测量噪声: {best_params['measurement_noise']}\n")
        f.write(f"  密度因子: {best_params['density_factor']}\n\n")
        f.write("样本评估结果:\n")
        f.write("--------------------------------\n")
        for result in results:
            f.write(f"样本 {result['sample_idx']}: MSE = {result['mse']:.4e}, 相关系数 = {result['correlation']:.4f}\n")
        
        # 计算平均指标
        avg_mse = np.mean([r['mse'] for r in results])
        avg_corr = np.mean([r['correlation'] for r in results])
        f.write("--------------------------------\n")
        f.write(f"平均 MSE: {avg_mse:.4e}\n")
        f.write(f"平均相关系数: {avg_corr:.4f}\n")
    
    print(f"所有样本处理完成，结果保存在 {result_dir} 目录")


def download_mdsplus_data(shot_num, server_path='192.168.20.11::/media/ennfusion/trees/exl50u'):
    """
    从MDSplus下载数据
    
    参数:
    - shot_num: 炮号
    - server_path: MDSplus服务器路径
    
    返回:
    - y_original: 原始相位信号
    - y_original_time: 原始相位时间
    - y_true: 真实密度信号
    - y_true_time: 真实密度时间
    """
    try:
        print(f"正在从MDSplus下载炮号 {shot_num} 的数据...")
        
        # 连接到MDSplus树
        tree = MDSplus.Tree('exl50u', shot_num, path=server_path)
        tree.setTimeContext(-1, 3, 1e-5)
        
        # 获取原始相位数据
        y_original = tree.getNode(r'\TEMP01').data()
        y_original_time = tree.getNode(r'\TEMP01').dim_of().data()
        
        # 获取真实密度数据
        y_true = tree.getNode(r'\HCN_NE001').data()
        y_true_time = tree.getNode(r'\HCN_NE001').dim_of().data()
        
        print(f"数据下载成功，相位数据长度: {len(y_original)}, 密度数据长度: {len(y_true)}")
        
        return y_original, y_original_time, y_true, y_true_time
    
    except Exception as e:
        print(f"下载数据时出错: {e}")
        raise


def plot_mdsplus_sample_detail(shot_num, process_noise, measurement_noise, density_factor, 
                              server_path='192.168.20.11::/media/ennfusion/trees/exl50u'):
    """
    绘制从MDSplus获取的单个样本的详细对比图
    
    参数:
    - shot_num: 炮号
    - process_noise: 卡尔曼滤波器过程噪声参数
    - measurement_noise: 卡尔曼滤波器测量噪声参数
    - density_factor: 密度计算的校正系数
    - server_path: MDSplus服务器路径
    
    返回:
    - mse: 均方根误差
    - corr: 相关系数
    """
    # 从MDSplus下载数据
    y_original, y_original_time, y_true, y_true_time = download_mdsplus_data(shot_num, server_path)
    
    # 确保长度一致
    min_length = min(len(y_original), len(y_true))
    y_original = y_original[:min_length]
    y_true = y_true[:min_length]
    
    # 使用指定参数处理
    unwrapped = kalman_unwrap(
        y_original, 
        process_noise=process_noise, 
        measurement_noise=measurement_noise
    )
    
    # 使用指定公式计算密度
    dl = 0
    y_processed = (unwrapped - dl) * 35.5 / 3.37 / 2 * density_factor
    
    # 计算评估指标
    mse = np.mean((y_processed - y_true) ** 2)
    corr = np.corrcoef(y_processed, y_true)[0, 1]
    
    # 绘制详细对比图
    plt.figure(figsize=(12, 8))
    
    # 上图：原始相位信号和解缠绕后的相位
    plt.subplot(2, 1, 1)
    plt.plot(y_original_time[:min_length], y_original, label='原始相位信号')
    plt.plot(y_original_time[:min_length], unwrapped, label='解缠绕后的相位')
    plt.title(f'炮号 {shot_num} 相位信号对比')
    plt.xlabel('时间 (s)')
    plt.ylabel('相位')
    plt.legend()
    plt.grid(True)
    
    # 下图：处理后的密度和真实密度
    plt.subplot(2, 1, 2)
    plt.plot(y_original_time[:min_length], y_processed, label='卡尔曼处理后')
    plt.plot(y_true_time[:min_length], y_true, 'r--', label='真实密度')
    plt.title(f'炮号 {shot_num} 密度对比: MSE = {mse:.4e}, 相关系数 = {corr:.4f}')
    plt.xlabel('时间 (s)')
    plt.ylabel('ne')
    plt.legend()
    plt.grid(True)
    
    plt.tight_layout()
    
    # 创建结果目录
    result_dir = 'e:\\ENN\\2025工作\\诊断-hcn拼接\\result_mdsplus'
    if not os.path.exists(result_dir):
        os.makedirs(result_dir)
    
    # 保存图像
    plt.savefig(f'{result_dir}\\kalman_shot_{shot_num}_detail.png')
    plt.close()  # 关闭图形，避免内存泄漏
    
    return mse, corr


def process_mdsplus_shots(shot_range=(8000, 9000), num_shots=10, best_params=None):
    """
    处理多个炮号的数据
    
    参数:
    - shot_range: 炮号范围，格式为(最小炮号, 最大炮号)
    - num_shots: 要处理的炮号数量
    - best_params: 最佳参数字典，如果为None则使用默认参数
    """
    # 如果未提供最佳参数，使用默认参数
    if best_params is None:
        best_params = {
            'process_noise': 0.01,
            'measurement_noise': 0.1,
            'density_factor': 1.0
        }
    
    # 创建结果目录
    result_dir = 'e:\\ENN\\2025工作\\诊断-hcn拼接\\result_mdsplus'
    if not os.path.exists(result_dir):
        os.makedirs(result_dir)
    
    # 随机选择炮号
    min_shot, max_shot = shot_range
    all_shots = list(range(min_shot, max_shot + 1))
    selected_shots = random.sample(all_shots, min(num_shots, len(all_shots)))
    
    print(f"已选择以下炮号进行处理: {selected_shots}")
    
    # 准备存储所有炮号的评估结果
    results = []
    
    # 循环处理所有选定的炮号
    for shot_num in selected_shots:
        try:
            print(f"处理炮号 {shot_num}...")
            mse, corr = plot_mdsplus_sample_detail(
                shot_num,
                best_params['process_noise'],
                best_params['measurement_noise'],
                best_params['density_factor']
            )
            results.append({
                'shot_num': shot_num,
                'mse': mse,
                'correlation': corr
            })
        except Exception as e:
            print(f"处理炮号 {shot_num} 时出错: {e}")
    
    # 生成汇总报告
    with open(f'{result_dir}\\mdsplus_summary_report.txt', 'w', encoding='utf-8') as f:
        f.write("卡尔曼滤波器相位解缠绕结果汇总 (MDSplus数据)\n")
        f.write("==========================================\n\n")
        f.write(f"使用参数:\n")
        f.write(f"  过程噪声: {best_params['process_noise']}\n")
        f.write(f"  测量噪声: {best_params['measurement_noise']}\n")
        f.write(f"  密度因子: {best_params['density_factor']}\n\n")
        f.write("炮号评估结果:\n")
        f.write("------------------------------------------\n")
        for result in results:
            f.write(f"炮号 {result['shot_num']}: MSE = {result['mse']:.4e}, 相关系数 = {result['correlation']:.4f}\n")
        
        # 计算平均指标
        if results:
            avg_mse = np.mean([r['mse'] for r in results])
            avg_corr = np.mean([r['correlation'] for r in results])
            f.write("------------------------------------------\n")
            f.write(f"平均 MSE: {avg_mse:.4e}\n")
            f.write(f"平均相关系数: {avg_corr:.4f}\n")
        else:
            f.write("没有成功处理的炮号数据\n")
    
    print(f"所有炮号处理完成，结果保存在 {result_dir} 目录")


if __name__ == "__main__":
    # 文件路径
    file_name = 'sjj1.mat'
    
    # 优化参数
    # print("正在优化卡尔曼滤波器参数...")
    # best_params = optimize_parameters(file_name)

    # 使用预设的默认参数
    best_params = {
        'process_noise': 0.001,
        'measurement_noise': 0.01,
        'density_factor': 0.8
    }
    print(f"使用默认参数: {best_params}")
    
    # 使用优化后的参数绘制对比图
    # print("绘制对比图...")
    # plot_comparison(
    #     file_name=file_name, 
    #     num_samples=4, 
    #     density_factor=best_params['density_factor']
    # )
    
    # 绘制所有样本的详细对比图
    # print("绘制所有样本的详细对比图...")
    # plot_all_samples(file_name, best_params)
    
    # 处理MDSplus数据
    print("处理MDSplus数据...")
    process_mdsplus_shots(
        shot_range=(8000, 9000),
        num_shots=10,
        best_params=best_params
    )
    
    print('处理完成!')